import os,cv2
import numpy as np

def findLoc(img):
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    _,threshImg = cv2.threshold(gray,5,255,cv2.THRESH_BINARY_INV)
    
    contours,hierarchy=cv2.findContours(threshImg,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    cnts = []
    for cnt in contours:
        cnts.append(list(cv2.boundingRect(cnt)))
    # print(cnts)
    return cnts[0]


def merge(src,dst):
    back = np.zeros((1080,1920,3),dtype=np.uint8)
    backTmp = np.zeros((1080,1920,3),dtype=np.uint8)

    backDownSize= np.zeros((540,1920,3),dtype=np.uint8)
    imgList = os.listdir(src)
    num = -1
    for i in range(len(imgList)):
        if(i<num):
            # back = np.zeros((1080,1920,3),dtype=np.uint8)
            # backTmp = np.zeros((1080,1920,3),dtype=np.uint8)
            continue
        back = np.zeros((1080,1920,3),dtype=np.uint8)
        backTmp = np.zeros((1080,1920,3),dtype=np.uint8)

        # num+=1
        file = imgList[i]
        print(file)
        img = cv2.imread(os.path.join(src,file))
        tmp_w,tmp_h = 0,0
        while(tmp_w<1920):
            _,w,h = img.shape[::-1]
            if(tmp_w+w<1920):
                back[slice(tmp_h,h+tmp_h),slice(tmp_w,tmp_w+w)] = img
                backTmp[slice(tmp_h,h+tmp_h),slice(tmp_w,tmp_w+w)] = (255,255,255)
                i+=1
                file = imgList[i]
                img = cv2.imread(os.path.join(src,file))
                if(tmp_w+w<1920):
                    tmp_w +=w
                # cv2.imwrite("/mnt/data/head/SCUT_HEAD_Part_B/merged/"+file.split(".")[0]+"_merged_up.jpg",back)

            else:
                break


        tmp_w = 0
        while(tmp_w<1920):
            _,w,h = img.shape[::-1]
            part = backTmp[slice(0,1080),slice(tmp_w,tmp_w+w)]

            startX,start_y,_,_ = findLoc(part)

            if(h+tmp_h+start_y<1080 and tmp_w+w<1920):
                back[slice(tmp_h+start_y,tmp_h+start_y+h),slice(tmp_w+startX,tmp_w+w+startX)] = img
                i+=1
                file = imgList[i]
                img = cv2.imread(os.path.join(src,file))
                if(tmp_w+w<1920):
                    tmp_w +=w
                print(tmp_h)
                cv2.imwrite("/mnt/data/head/SCUT_HEAD_Part_B/merged/"+file.split(".")[0]+"_merged_down.jpg",back)

            else:

                break

        cv2.imwrite("/mnt/data/head/SCUT_HEAD_Part_B/merged/"+file.split(".")[0]+"_merged_down.jpg",back)
        num = i



def main():
    src = "/mnt/data/head/SCUT_HEAD_Part_B/JPEGImages"
    dst = "/mnt/data/head/SCUT_HEAD_Part_B/merged"
    merge(src,dst)

if __name__ == "__main__":
    main()